import torch
import torch.nn as nn
import torch.nn.functional as F

from .base_laysers import Encoder, Q_Net


class AttModel(nn.Module):
    def __init__(self, n_node, din, hidden_dim, dout):
        super(AttModel, self).__init__()
        self.fcv = nn.Linear(din, hidden_dim)
        self.fck = nn.Linear(din, hidden_dim)
        self.fcq = nn.Linear(din, hidden_dim)
        self.fcout = nn.Linear(hidden_dim, dout)

    def forward(self, x, mask: torch.Tensor):
        mask = mask.to(dtype=torch.int8)
        v = F.relu(self.fcv(x))
        q = F.relu(self.fcq(x))
        k = F.relu(self.fck(x)).permute(0, 2, 1)
        att = F.softmax(torch.mul(torch.bmm(q, k), mask) - 9e15 * (1 - mask), dim=2)

        out = torch.bmm(att, v)
        out = F.relu(self.fcout(out))
        return out


class DGN_R1(nn.Module):
    def __init__(self, n_agent, num_inputs, hidden_dim, num_actions):
        super(DGN_R1, self).__init__()

        self.encoder = Encoder(num_inputs, hidden_dim)
        self.att_1 = AttModel(n_agent, hidden_dim, hidden_dim, hidden_dim)
        self.q_net = Q_Net(hidden_dim, num_actions)

    def forward(self, x, mask):
        h1 = self.encoder(x)
        h2 = self.att_1(h1, mask)
        q = self.q_net(h2)
        return q, mask


class DGN_R2(nn.Module):
    def __init__(self, n_agent, num_inputs, hidden_dim, num_actions):
        super(DGN_R2, self).__init__()

        self.encoder = Encoder(num_inputs, hidden_dim)
        self.layer_num = 4
        self.layers = nn.ModuleList(
            [AttModel(n_agent, hidden_dim, hidden_dim, hidden_dim) for _ in range(self.layer_num)])
        self.q_net = Q_Net(hidden_dim * 2, num_actions)

    def forward(self, x: torch.Tensor, mask):
        batch_size, n_agent, num_dims, dim_length = x.shape
        h1 = torch.reshape(x, (batch_size, n_agent, num_dims * dim_length))

        h1 = self.encoder(h1)

        temp1 = h1
        temp2 = None
        for layer in self.layers:
            temp2 = layer(temp1, mask)
            temp1 = temp2

        q = self.q_net(torch.cat([h1, temp2], dim=-1))
        return q, mask
